#!/usr/bin/env python
# coding: utf-8
import os
# 强制使用 torch 后端
os.environ["KERAS_BACKEND"] = "torch"

import keras
from keras import Sequential
from keras.layers import LSTM, Dropout, Dense
from keras.optimizers import Adam
from keras.metrics import AUC, Recall
from keras.callbacks import EarlyStopping
from tqdm.keras import TqdmCallback
from utils.landslide_utils import predict_class_prob, f1_loss, F1Score
from utils.general_utils import evaluate_metrics_plotROC


def lstm_model(input_shape=(1, 18),
               nb_mid_layers=5,
               nb_class=2,
               node_num=10,
               lr_rate=0.001,
               loss_fun='categorical_crossentropy'):
    """
    创建LSTM模型

    Parameters:
    -----------
    input_shape : tuple
        输入形状，默认为(1, 18)
    nb_mid_layers : int
        中间LSTM层数，默认为5
    nb_class : int
        分类数量，默认为2（二分类）
    node_num : int
        LSTM单元数，默认为10
    lr_rate : float
        学习率，默认为0.001
    loss_fun : str
        损失函数，默认为'categorical_crossentropy'

    Returns:
    --------
    model : keras.Model
        编译好的Keras模型
    """
    model = Sequential()
    # 添加一个独立的输入层
    model.add(keras.layers.InputLayer(shape=input_shape))
    model.add(LSTM(node_num, return_sequences=True))

    for i in range(nb_mid_layers):
        model.add(LSTM(node_num, return_sequences=True))
        model.add(Dropout(0.5))

    model.add(LSTM(node_num))
    model.add(Dropout(0.5))
    model.add(Dense(nb_class, activation='softmax'))

    # 使用自定义的F1Score指标
    model.compile(
        Adam(learning_rate=lr_rate),
        loss=f1_loss,
        metrics=[AUC(), F1Score(), Recall()]
    )

    return model


def fit(x_train,
        y_train,
        nb_epochs=200,
        verbose=0,
        input_shape=(1, 18),
        nb_mid_layers=5,
        node_num=64,
        lr_rate=0.001,
        class_weight=None,
        bat_size=1000,
        val_split=0.2,
        val_data=None,
        use_early_stop=False):
    """
    训练LSTM模型

    Parameters:
    -----------
    x_train : array-like
        训练数据特征
    y_train : array-like
        训练数据标签
    nb_epochs : int
        训练轮数，默认为200
    verbose : int
        详细程度，默认为0（不显示进度条）
    input_shape : tuple
        输入形状，默认为(1, 18)
    nb_mid_layers : int
        中间LSTM层数，默认为5
    node_num : int
        LSTM单元数，默认为64
    lr_rate : float
        学习率，默认为0.001
    class_weight : dict
        类别权重，用于处理不平衡数据
    bat_size : int
        批次大小，默认为1000
    val_split : float
        验证集分割比例，默认为0.2
    val_data : tuple
        验证数据，格式为(x_val, y_val_onehot)
    use_early_stop : keras.callbacks.EarlyStopping
        是否使用早停。默认为False，不使用早停。

    Returns:
    --------
    model : keras.Model
        训练好的模型
    history : keras.History
        训练历史记录
    """
    model = lstm_model(
        input_shape=input_shape,
        nb_mid_layers=nb_mid_layers,
        nb_class=2,
        node_num=node_num,
        lr_rate=lr_rate
    )
    custom_callbacks = [TqdmCallback(verbose=verbose)]
    if use_early_stop:
        # 使用F1Score作为监控指标,Keras无法自动确定自定义指标（如F1Score）应该是最大化还是最小化，因此需要显式指定 mode参数，不能使用auto了。
        early_stop = EarlyStopping(monitor='val_f1', patience=50, verbose=2, mode='max')
        custom_callbacks = [TqdmCallback(verbose=verbose), early_stop]

    history = model.fit(
        x_train,
        y_train,
        validation_split=val_split,
        validation_data=val_data,
        shuffle=True,
        batch_size=bat_size,
        epochs=nb_epochs,
        verbose=verbose,
        class_weight=class_weight,
        callbacks=custom_callbacks
    )

    return model, history


def run_once_lstm(x_train,
                  y_train_onehot,
                  x_test,
                  y_test_onehot,
                  y_test_not_onehot,
                  channels,
                  nb_epochs,
                  nb_mid_layers,
                  cls_weight,
                  nb_pca_or_nb_features,
                  return_pred_y=False,
                  verbose=0,
                  node_num=64,
                  lr_rate=0.0001,
                  bat_size=1000,
                  val_split=0.2,
                  val_data=None,
                  show_roc=False,
                  print_results=True,
                  use_early_stop=False,
                  ):
    title = f'lstm train {nb_epochs} epochs'
    model, history = fit(x_train, y_train_onehot,
                         nb_epochs=nb_epochs,
                         verbose=verbose,
                         input_shape=(channels, nb_pca_or_nb_features),
                         nb_mid_layers=nb_mid_layers,
                         class_weight=cls_weight,
                         node_num=node_num,
                         lr_rate=lr_rate,
                         bat_size=bat_size,
                         val_split=val_split,
                         val_data = val_data,
                         use_early_stop=use_early_stop)
    y_pred_class, y_pred_proba = predict_class_prob(model, x_test, bat_size)

    # 注意：这里的y_test_r为1d，不能使用y_test, y_test经过了独热编码为2D
    n = len(y_test_not_onehot)
    k = x_test.shape[1]
    return_AIC_params = {
        'return_AIC': True,
        'n': n,
        'k': k
    }
    metrics_dict = evaluate_metrics_plotROC(y_test_not_onehot, y_pred_class, y_pred_proba,
                                            title=title, show_roc=show_roc,
                                            print_results=print_results,
                                            return_AIC_params=return_AIC_params)

    if return_pred_y:
        return [y_pred_class, y_pred_proba], model, history
    else:
        return metrics_dict, model, history
